#ifndef UnstructMG_ITER_SOLVER_HPP
#define UnstructMG_ITER_SOLVER_HPP

#include "operator.hpp"

typedef enum {PREC, OPER, AXPY, DOT, NUM_KRYLOV_RECORD} KRYLOV_RECORD_TYPES; 

// 虚基类IterativeSolver，支持两种精度
// 继承自虚基类->Solver->Operator，重写了SetOperator()，还需下一层子类重写Mult_impl()
// 预条件的存储类型pc_data_t, 预条件的计算类型pc_calc_t, 外迭代的精度ksp_t
template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
class IterativeSolver : public Solver<idx_t, ksp_t, ksp_t, dof> {
protected:
    // oper(外迭代的算子/矩阵)，可以和prec(预条件子)采用不一样的精度
    const Operator<idx_t, ksp_t, dof> *oper = nullptr;
    // 预条件子的存储采用低精度，计算采用高精度
    Solver<idx_t, ksp_t, pc_calc_t, dof> *prec = nullptr;
public:
    int max_iter = 10, print_level = -1;
    double rel_tol = 0.0, abs_tol = 0.0;// 用高精度的收敛判断

    // stats
    mutable int final_iter = 0, converged = 0;
    mutable double final_norm;
    mutable double part_times[NUM_KRYLOV_RECORD];

    IterativeSolver();
    ~IterativeSolver();

    void SetRelTol(double rtol);
    void SetAbsTol(double atol);
    void SetMaxIter(int max_it);
    void SetPrintLevel(int print_lvl);
    /// This should be called before SetOperator
    virtual void SetPreconditioner(Solver<idx_t, ksp_t, pc_calc_t, dof> & pr);
    /// Also calls SetOperator for the preconditioner if there is one
    virtual void SetOperator(const Operator<idx_t, ksp_t, dof> & op);
    
    int GetNumIterations() const;
    int GetConverged() const;
    double GetFinalNorm() const;

    virtual void truncate();

    // 迭代法里的点积（默认采用双精度）
    double Dot(const par_Vector<idx_t, ksp_t, dof> & x, const par_Vector<idx_t, ksp_t, dof> & y) const;
    double Norm(const par_Vector<idx_t, ksp_t, dof> & x) const;


};

typedef IterativeSolver<int, double, double, double, 1> IterSolver_I32All64_DOF1;
typedef IterativeSolver<int, double, double, double, 2> IterSolver_I32All64_DOF2;
typedef IterativeSolver<int, double, double, double, 3> IterSolver_I32All64_DOF3;
typedef IterativeSolver<int, double, double, double, 4> IterSolver_I32All64_DOF4;

typedef IterativeSolver<int, float , double, double, 1> IterSolver_I32K64P64D32_DOF1;
typedef IterativeSolver<int, float , double, double, 2> IterSolver_I32K64P64D32_DOF2;
typedef IterativeSolver<int, float , double, double, 3> IterSolver_I32K64P64D32_DOF3;
typedef IterativeSolver<int, float , double, double, 4> IterSolver_I32K64P64D32_DOF4;

typedef IterativeSolver<int, float , float , double, 1> IterSolver_I32K64P32D32_DOF1;
typedef IterativeSolver<int, float , float , double, 2> IterSolver_I32K64P32D32_DOF2;
typedef IterativeSolver<int, float , float , double, 3> IterSolver_I32K64P32D32_DOF3;
typedef IterativeSolver<int, float , float , double, 4> IterSolver_I32K64P32D32_DOF4;

typedef IterativeSolver<int, float , float , float , 1> IterSolver_I32All32_DOF1;
typedef IterativeSolver<int, float , float , float , 2> IterSolver_I32All32_DOF2;
typedef IterativeSolver<int, float , float , float , 3> IterSolver_I32All32_DOF3;
typedef IterativeSolver<int, float , float , float , 4> IterSolver_I32All32_DOF4;

#ifdef USE_FP16
typedef IterativeSolver<int, __fp16, float , double, 1> IterSolver_I32K64P32D16_DOF1;
typedef IterativeSolver<int, __fp16, float , double, 2> IterSolver_I32K64P32D16_DOF2;
typedef IterativeSolver<int, __fp16, float , double, 3> IterSolver_I32K64P32D16_DOF3;
typedef IterativeSolver<int, __fp16, float , double, 4> IterSolver_I32K64P32D16_DOF4;

typedef IterativeSolver<int, __fp16, float , float , 1> IterSolver_I32K32P32D16_DOF1;
typedef IterativeSolver<int, __fp16, float , float , 2> IterSolver_I32K32P32D16_DOF2;
typedef IterativeSolver<int, __fp16, float , float , 3> IterSolver_I32K32P32D16_DOF3;
typedef IterativeSolver<int, __fp16, float , float , 4> IterSolver_I32K32P32D16_DOF4;
#endif

template<typename idx_t, typename src_t, typename dst_t, int dof>
void VecTransPrec(const par_Vector<idx_t, src_t, dof> & src_vec, par_Vector<idx_t, dst_t, dof> & dst_vec);
extern template void VecTransPrec(const par_Vector_I32F64_DOF1 &, par_Vector_I32F64_DOF1 &);
extern template void VecTransPrec(const par_Vector_I32F64_DOF2 &, par_Vector_I32F64_DOF2 &);
extern template void VecTransPrec(const par_Vector_I32F64_DOF3 &, par_Vector_I32F64_DOF3 &);
extern template void VecTransPrec(const par_Vector_I32F64_DOF4 &, par_Vector_I32F64_DOF4 &);

extern template void VecTransPrec(const par_Vector_I32F64_DOF1 &, par_Vector_I32F32_DOF1 &);
extern template void VecTransPrec(const par_Vector_I32F64_DOF2 &, par_Vector_I32F32_DOF2 &);
extern template void VecTransPrec(const par_Vector_I32F64_DOF3 &, par_Vector_I32F32_DOF3 &);
extern template void VecTransPrec(const par_Vector_I32F64_DOF4 &, par_Vector_I32F32_DOF4 &);

extern template void VecTransPrec(const par_Vector_I32F32_DOF1 &, par_Vector_I32F32_DOF1 &);
extern template void VecTransPrec(const par_Vector_I32F32_DOF2 &, par_Vector_I32F32_DOF2 &);
extern template void VecTransPrec(const par_Vector_I32F32_DOF3 &, par_Vector_I32F32_DOF3 &);
extern template void VecTransPrec(const par_Vector_I32F32_DOF4 &, par_Vector_I32F32_DOF4 &);

extern template void VecTransPrec(const par_Vector_I32F32_DOF1 &, par_Vector_I32F64_DOF1 &);
extern template void VecTransPrec(const par_Vector_I32F32_DOF2 &, par_Vector_I32F64_DOF2 &);
extern template void VecTransPrec(const par_Vector_I32F32_DOF3 &, par_Vector_I32F64_DOF3 &);
extern template void VecTransPrec(const par_Vector_I32F32_DOF4 &, par_Vector_I32F64_DOF4 &);

template<typename idx_t, typename data_t, typename res_t, int dof>
res_t vec_dot_intra_proc(const par_Vector<idx_t, data_t, dof> & x, const par_Vector<idx_t, data_t, dof> & y);
extern template double vec_dot_intra_proc(const par_Vector_I32F32_DOF1 &, const par_Vector_I32F32_DOF1 &);
extern template double vec_dot_intra_proc(const par_Vector_I32F32_DOF2 &, const par_Vector_I32F32_DOF2 &);
extern template double vec_dot_intra_proc(const par_Vector_I32F32_DOF3 &, const par_Vector_I32F32_DOF3 &);
extern template double vec_dot_intra_proc(const par_Vector_I32F32_DOF4 &, const par_Vector_I32F32_DOF4 &);
extern template double vec_dot_intra_proc(const par_Vector_I32F64_DOF1 &, const par_Vector_I32F64_DOF1 &);
extern template double vec_dot_intra_proc(const par_Vector_I32F64_DOF2 &, const par_Vector_I32F64_DOF2 &);
extern template double vec_dot_intra_proc(const par_Vector_I32F64_DOF3 &, const par_Vector_I32F64_DOF3 &);
extern template double vec_dot_intra_proc(const par_Vector_I32F64_DOF4 &, const par_Vector_I32F64_DOF4 &);


#endif